from math import log, exp


class Prob:
    def __init__(self):
        self.data = dict()
        self.sum = 0.0

    def getCount(self, key):
        if key in self.data:
            return self.data[key]
        return None

    def getProb(self, key):
        count = self.getCount(key)
        if count is None:
            return 0
        return float(count) / self.sum

    def add(self, key, val):
        self.sum += val
        if key not in self.data:
            self.data[key] = val
        else:
            self.data[key] += val


class NaiveBayes:
    def __init__(self):
        self.data = dict()
        self.sum = 0.0

    def train(self, dataSet):
        for item in dataSet:
            label = item[1]
            if label not in self.data:
                self.data[label] = Prob()
            for word in item[0]:
                self.data[label].add(word, 1)
        self.sum = sum(map(lambda x: self.data[x].sum, self.data.keys()))

    def predict(self, item):
        tmp = dict()
        for key in self.data:
            tmp[key] = log(self.data[key].sum) - log(self.sum)
            for word in item:
                tmp[key] += log(self.data[key].getProb(word))
        result = 0
        prob = 0
        for key1 in self.data:
            now = 0
            try:
                for key2 in self.data:
                    now += exp(tmp[key2] - tmp[key1])
                    now = 1 / now
            except OverflowError:
                now = 0
            if now > prob:
                result = key1
                prob = now
        return result, prob

